from .config_manager import ConfigManager
from .device_manager import DeviceManager
from .model_manager import ModelManager

class MiniCPMVDeploy:
    """
    MiniCPM-V 部署类，整合配置管理、设备管理和模型管理功能
    """
    def __init__(self, config_file=None):
        # 初始化配置管理器
        self.config_manager = ConfigManager(config_file)
        self.config = self.config_manager.config
        
        # 初始化设备管理器
        self.device_manager = DeviceManager(self.config)
        self.device = self.device_manager.get_device()
        self.dtype = self.device_manager.get_dtype()
        
        # 初始化模型管理器
        self.model_manager = ModelManager(self.config)
        
        print(f"\n===== MiniCPM-V 部署初始化完成 =====")
        print(f"配置文件: {self.config_manager.config_file}")
        print(f"设备: {self.device}")
        print(f"数据类型: {self.dtype}")
        print(f"模型目录: {self.config['paths']['model_dir']}")
        print(f"输出目录: {self.config['paths']['output_dir']}")
        print(f"==================================\n")
    
    def load_model(self):
        """
        加载模型
        """
        return self.model_manager.load_model(self.device, self.dtype)
    
    def get_model(self):
        """
        获取加载的模型
        """
        return self.model_manager.get_model()
    
    def get_tokenizer(self):
        """
        获取加载的tokenizer
        """
        return self.model_manager.get_tokenizer()
    
    def update_config(self, updates):
        """
        更新配置
        """
        success = self.config_manager.update_config(updates)
        if success:
            self.config = self.config_manager.config
            # 重新初始化设备管理器以应用新配置
            self.device_manager = DeviceManager(self.config)
            self.device = self.device_manager.get_device()
            self.dtype = self.device_manager.get_dtype()
        return success
    
    def get_config(self):
        """
        获取当前配置
        """
        return self.config
    
    def unload_model(self):
        """
        卸载模型，释放内存
        """
        self.model_manager.unload_model()
    
    def get_device_info(self):
        """
        获取设备信息
        """
        if self.device_manager.is_cuda_available():
            return self.device_manager.get_gpu_info()
        else:
            return {
                "device": self.device,
                "dtype": self.dtype
            }
    
    def is_amp_enabled(self):
        """
        检查是否启用自动混合精度
        """
        return self.device_manager.is_amp_enabled()